import numpy as np


pc = np.array([
    [10,10,0],
    [1,10,0],
    [10,1,0],
    [0.9,0.9,0],
    [1.1,0.9,0],
    [0,0.9,0.98],
    [100,100.9,0.98],
    [101.01,100.9,0.98],
])

bboxes = np.array([
    [0,0,0,2,2,2,0],
    [100,100,0,2,2,2,0],
])


# 定义边界框角点生成的函数
def get_bbox_corners(center, size, heading_angle):
    # 计算旋转矩阵
    R = np.array([
        [np.cos(heading_angle), -np.sin(heading_angle), 0],
        [np.sin(heading_angle), np.cos(heading_angle), 0],
        [0, 0, 1]
    ])

    # 定义八个角点的偏移量
    half_size = size / 2.0
    corners_offset = np.array([
        [-half_size[0], -half_size[1], -half_size[2]],
        [half_size[0], -half_size[1], -half_size[2]],
        [half_size[0], half_size[1], -half_size[2]],
        [-half_size[0], half_size[1], -half_size[2]],
        [-half_size[0], -half_size[1], half_size[2]],
        [half_size[0], -half_size[1], half_size[2]],
        [half_size[0], half_size[1], half_size[2]],
        [-half_size[0], half_size[1], half_size[2]]
    ])

    # 旋转角点偏移量
    corners_3d = (R @ corners_offset.T).T

    # 加上中心点坐标
    corners_3d += center

    return corners_3d


# 计算所有边界框的角点
bbox_corners = np.array([get_bbox_corners(bbox[:3], bbox[3:6], bbox[6]) for bbox in bboxes])


# 计算每个点是否在边界框内（通过构造AABB包围盒快速检查）
def is_point_in_aabb(point, min_corner, max_corner):
    in_box = (point[None,:,:]>=min_corner[:,None,:]) & (point[None,:,:]<=max_corner[:,None,:])
    in_box =  np.all(in_box,axis=-1)
    return in_box


# 扩展边界框角点，方便进行点云内点判断
min_corners = bbox_corners.min(axis=1)  # [10, 3]
max_corners = bbox_corners.max(axis=1)  # [10, 3]


# 使用numpy的向量化比较来判断点是否在AABB包围盒内
point_mask = is_point_in_aabb(pc, min_corners, max_corners)  # [10000, 10, 1]

# 计数每个边界框内的点数
point_counts = np.sum(point_mask, axis=1)

print(point_counts)